#include "rpc.h"
#include "helper.h"
#include <string.h>
#include <stdlib.h>
#include <stdio.h>
#include <signal.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#define NONBLOCKING
#define MAX_DATA2_LEN 2000000

struct rpc_server {
    /* Add variable(s) for server state */
    PLNode head;
    int id;
    int port;
};

static int server_kill = 0;

static void handle_sigint(int sig, siginfo_t *si, void *ctx) {
    server_kill = 1;
}

rpc_server *rpc_init_server(int port) {
    // init state
    rpc_server* res = (rpc_server*)malloc(sizeof(rpc_server));
    res->head = create_dummy_head();
    res->id = 0;
    res->port = port;

    // register a kill signal handler
    struct sigaction act;
    act.sa_sigaction = handle_sigint;
    act.sa_flags = SA_SIGINFO;
    sigemptyset(&act.sa_mask);
    if (sigaction(SIGINT, &act, NULL) == -1) {
        fprintf(stderr, "Couldn't set up SIGINT handler: %s\n", strerror(errno));
        exit(EXIT_FAILURE);
    }

    return res;
}

int rpc_register(rpc_server *srv, char *name, rpc_handler handler) {
    if(valid_name(name) == 0 || strcmp(name, "rpc_find") == 0) {
        return -1;
    }
    int id = srv->id++;
    node_register(srv->head, id, name, handler);
    
    return id;
}

static void rpc_close(rpc_server *srv, int conn_fd, int server_sockfd) {
    if(srv == NULL) {
        return;
    }
    if(srv->head != NULL) {
        free_list(srv->head);
    }
    free(srv);
    if(conn_fd > 0)
        close(conn_fd);
    if(server_sockfd > 0)
        close(server_sockfd);
}

// exit with error flag set
static inline void error_exit(char* str, rpc_server *srv, int conn_fd, int server_sockfd) {
    rpc_close(srv, conn_fd, server_sockfd);
    perror(str);
    exit(EXIT_FAILURE);
}
// just exit
static inline void child_exit(rpc_server *srv, int conn_fd, int server_sockfd) {
    rpc_close(srv, conn_fd, server_sockfd);
    exit(EXIT_SUCCESS);
}

// reciver
void rpc_serve_all(rpc_server *srv) {
    if(srv == NULL) {
        return;
    }
    // open and listen a socket
    int server_sockfd = socket(AF_INET6,SOCK_STREAM, 0);

    struct sockaddr_in6 server_sockaddr, client_sockaddr;
    memset((void*)&server_sockaddr, 0, sizeof(struct sockaddr_in6));
    server_sockaddr.sin6_addr = in6addr_any;
    server_sockaddr.sin6_family = AF_INET6;
    server_sockaddr.sin6_port = htons(srv->port);

    int enable = 1;
    if (setsockopt(server_sockfd, SOL_SOCKET, SO_REUSEADDR, &enable, sizeof(int)) < 0) {
        error_exit("setsockopt", srv, -1, -1);
    }

    if(bind(server_sockfd, (struct sockaddr *)&server_sockaddr, sizeof(server_sockaddr)) == -1) {
        error_exit("bind", srv, -1, -1);
    }

    if(listen(server_sockfd, 10) == -1) {
        error_exit("listen", srv, -1, -1);
    }

    while(1) {
        uint32_t length;
        // for now, serve as many client as we can
        int conn_fd = accept(server_sockfd, (struct sockaddr*)&client_sockaddr, &length);
        if(conn_fd < 0) {
            error_exit("accept", srv, -1, server_sockfd);
        }

        if(fork() == 0) {
            int down = 0;
            while(!down) {
                if(server_kill) {
                    error_exit("server killed!\n", srv, conn_fd, server_sockfd);
                }

                int type;
                uint32_t data2_len;
                rpc_data* arg_data, *resp_data;
                PLNode node;
                int name_length;
                char name_buff[10000]; // TODO: use malloc

                // |TYPE|name_length|name|rpc_data|
                if(recv(conn_fd, &type, sizeof(int), 0) < 0) {
                    child_exit(srv, conn_fd, server_sockfd);
                }
                type = ntohl(type);

                if(type >= NONE && type <= CLOSE) {
                    switch (type)
                    {
                        case REQ:
                            if(recv(conn_fd, &name_length, sizeof(int), 0) < 0) {
                                child_exit(srv, conn_fd, server_sockfd);
                            }
                            name_length = ntohl(name_length);
                            if(recv(conn_fd, name_buff, name_length, 0) < 0) {
                                child_exit(srv, conn_fd, server_sockfd);
                            }

                            arg_data = (rpc_data*)malloc(sizeof(rpc_data));
                            memset(arg_data, 0, sizeof(rpc_data));
                            if(recv(conn_fd, &(arg_data->data1), sizeof(arg_data->data1), 0) < 0) {
                                rpc_data_free(arg_data);
                                child_exit(srv, conn_fd, server_sockfd);
                            }
                            if(recv(conn_fd, &data2_len, sizeof(data2_len), 0) < 0) {
                                rpc_data_free(arg_data);
                                child_exit(srv, conn_fd, server_sockfd);
                            }
                            arg_data->data1 = ntohl(arg_data->data1);
                            arg_data->data2_len = ntohl(data2_len);
                            if(arg_data->data2_len) {
                                arg_data->data2 = malloc(arg_data->data2_len);
                                if(recv(conn_fd, arg_data->data2, arg_data->data2_len, 0) < 0) {
                                    rpc_data_free(arg_data);
                                    child_exit(srv, conn_fd, server_sockfd);
                                }
                            }

                            node = find_node_by_name(srv->head, name_buff);
                            if(strcmp(name_buff, "rpc_find") == 0) {
                                // refind the function using name in arg_data->data2
                                node = find_node_by_name(srv->head, (char*)arg_data->data2);
                                // send response back to client
                                // |TYPE|rpc_data|
                                type = htonl(RESP);
                                int result = htonl(node != NULL);
                                if(send(conn_fd, &type, sizeof(int), 0) < 0) {
                                    rpc_data_free(arg_data);
                                    child_exit(srv, conn_fd, server_sockfd);
                                }
                                if(send(conn_fd, &result, sizeof(arg_data->data1), 0) < 0) { // data1
                                    rpc_data_free(arg_data);
                                    child_exit(srv, conn_fd, server_sockfd);
                                }
                                result = 0;
                                if(send(conn_fd, &result, sizeof(uint32_t), 0) < 0) { // data2_length
                                    rpc_data_free(arg_data);
                                    child_exit(srv, conn_fd, server_sockfd);
                                }
                                rpc_data_free(arg_data);
                            }else if(node) {
                                resp_data = (node->handler)(arg_data);
                                // send response back to client
                                // |TYPE|rpc_data|
                                if(resp_data) {
                                    type = htonl(RESP);
                                    resp_data->data1 = htonl(resp_data->data1);
                                    data2_len = htonl(resp_data->data2_len);
                                    if(send(conn_fd, &type, sizeof(int), 0) < 0) { // type
                                        rpc_data_free(resp_data);
                                        rpc_data_free(arg_data);
                                        child_exit(srv, conn_fd, server_sockfd);
                                    }
                                    if(send(conn_fd, &(resp_data->data1), sizeof(resp_data->data1), 0) < 0) { // data1
                                        rpc_data_free(resp_data);
                                        rpc_data_free(arg_data);
                                        child_exit(srv, conn_fd, server_sockfd);
                                    }
                                    if(send(conn_fd, &data2_len, sizeof(data2_len), 0) < 0) { // data2_len
                                        rpc_data_free(resp_data);
                                        rpc_data_free(arg_data);
                                        child_exit(srv, conn_fd, server_sockfd);
                                    }
                                    if(resp_data->data2_len) {
                                        if(send(conn_fd, resp_data->data2, resp_data->data2_len, 0) < 0) { // data2
                                            rpc_data_free(resp_data);
                                            rpc_data_free(arg_data);
                                            child_exit(srv, conn_fd, server_sockfd);
                                        }
                                    }
                                }else {
                                    type = htonl(ERROR);
                                    if(send(conn_fd, &type, sizeof(int), 0) < 0) {
                                        rpc_data_free(resp_data);
                                        rpc_data_free(arg_data);
                                        child_exit(srv, conn_fd, server_sockfd);
                                    }
                                }
                                // free arg and resp data
                                rpc_data_free(resp_data);
                                rpc_data_free(arg_data);
                            }else {
                                // not found
                                type = htonl(ERROR);
                                rpc_data_free(arg_data);
                                if(send(conn_fd, &type, sizeof(int), 0) < 0) {
                                    child_exit(srv, conn_fd, server_sockfd);
                                }
                            }
                            break;
                        case CLOSE:
                            down = 1;
                            break;
                        default:
                            break;
                    }
                }else {
                    // undefined type
                    child_exit(srv, conn_fd, server_sockfd);
                }
            }
            child_exit(srv, conn_fd, server_sockfd);
        }
        close(conn_fd);
    }
}

struct rpc_client {
    /* Add variable(s) for client state */
    char* addr;
    int port;
};

struct rpc_handle {
    /* Add variable(s) for handle */
    char* name;
};

rpc_client *rpc_init_client(char *addr, int port) {
    rpc_client* result = (rpc_client*)malloc(sizeof(rpc_client));

    result->addr = (char*)malloc(strlen(addr) + 1);
    strcpy(result->addr, addr);
    result->port = port;
    
    return result;
}

static int client_connet(rpc_client *cl) {
    int sockfd = socket(AF_INET6, SOCK_STREAM, 0);
    struct sockaddr_in6 servaddr;
    memset(&servaddr, 0, sizeof(servaddr));
    servaddr.sin6_family = AF_INET6;
    servaddr.sin6_port = htons(cl->port);
    inet_pton(AF_INET6, cl->addr, &(servaddr.sin6_addr));

    if (connect(sockfd, (struct sockaddr *)&servaddr, sizeof(servaddr)) < 0) {
        perror("connect");
        exit(EXIT_FAILURE);
    }
    return sockfd;
}

// |TYPE|name_length|name|rpc_data|
static int client_send_req(rpc_client *cl, char* func_name, int conn_fd, rpc_data* data) {
    int type = htonl(REQ);
    int name_length = htonl(strlen(func_name) + 1);

    if(send(conn_fd, &type, sizeof(int), 0) < 0) // type
        return -1;
    if(send(conn_fd, &name_length, sizeof(int), 0) < 0) // name_length
        return -1;
    if(send(conn_fd, func_name, ntohl(name_length), 0) < 0) // name
        return -1;

    int data1 = htonl(data->data1);
    uint32_t data2_length = htonl(data->data2_len);
    if(send(conn_fd, &data1, sizeof(data->data1), 0) < 0) // rpc_data->data1
        return -1;
    if(send(conn_fd, &data2_length, sizeof(data2_length), 0) < 0) // rpc_data->data2_len
        return -1;
    if(send(conn_fd, data->data2, data->data2_len, 0) < 0) // rpc_data->data2
        return -1;

    return 0; // valid
}

static void client_send_close(rpc_client *cl, int conn_fd) {
    int type = htonl(CLOSE);
    send(conn_fd, &type, sizeof(int), 0); // type
}

// |TYPE|rpc_data|
static rpc_data* client_get_resp(rpc_client *cl, int conn_fd) {
    rpc_data* res = NULL;
    int type;
    uint32_t data2_len;
    if(recv(conn_fd, &type, sizeof(int), 0) < 0)
        return NULL;
    type = ntohl(type);
    if(type == RESP) {
        res = (rpc_data*)malloc(sizeof(rpc_data));
        memset(res, 0, sizeof(rpc_data));
        if(recv(conn_fd, &(res->data1), sizeof(res->data1), 0) < 0) {
            free(res);
            return NULL;
        }
        if(recv(conn_fd, &(data2_len), sizeof(data2_len), 0) < 0) {
            free(res);
            return NULL;
        }
        res->data1 = ntohl(res->data1);
        res->data2_len = ntohl(data2_len);
        if(res->data2_len) {
            res->data2 = malloc(res->data2_len);
            if(recv(conn_fd, res->data2, res->data2_len, 0) < 0) {
                free(res->data2);
                free(res);
                return NULL;
            }
        }
    }else if(type == ERROR) {
        return NULL;
    }
    return res;
}

rpc_handle *rpc_find(rpc_client *cl, char *name) {
    if(cl == NULL || name == NULL || !valid_name(name)) {
        return NULL;
    }
    int conn_fd = client_connet(cl);
    if(conn_fd < 0) {
        return NULL;
    }
    rpc_handle* res = NULL;
    rpc_data* resp_data = NULL;

    rpc_data* arg_data = (rpc_data*)malloc(sizeof(rpc_data));
    arg_data->data1 = 0; // dontcare
    arg_data->data2_len = strlen(name) + 1;
    arg_data->data2 = (char*)malloc(arg_data->data2_len);
    strcpy(arg_data->data2, name);

    if(client_send_req(cl, "rpc_find", conn_fd, arg_data) < 0) {
        rpc_data_free(arg_data);
        close(conn_fd);
        return NULL;
    }
    resp_data = client_get_resp(cl, conn_fd);
    if(resp_data && resp_data->data1) {
        res = (rpc_handle*)malloc(sizeof(rpc_handle));
        res->name = name;
    }
    client_send_close(cl, conn_fd);

    rpc_data_free(resp_data);
    rpc_data_free(arg_data);
    close(conn_fd);
    return res;
}

rpc_data *rpc_call(rpc_client *cl, rpc_handle *h, rpc_data *payload) {
    if(cl == NULL || h == NULL || payload == NULL) {
        return NULL;
    }
    int conn_fd = client_connet(cl);
    if(conn_fd < 0) {
        return NULL;
    }
    rpc_data* resp_data = NULL;

    if(client_send_req(cl, h->name, conn_fd, payload) < 0) {
        close(conn_fd);
        return NULL;
    }
    resp_data = client_get_resp(cl, conn_fd);

    client_send_close(cl, conn_fd);
    close(conn_fd);
    return resp_data;
}

void rpc_close_client(rpc_client *cl) {
    if(cl->addr)
        free(cl->addr);
    free(cl);
}

void rpc_data_free(rpc_data *data) {
    if (data == NULL) {
        return;
    }
    if (data->data2 != NULL) {
        free(data->data2);
    }
    free(data);
}
